﻿#define NDNS_INTERNAL

using Ndns.Packet.Abstractions;
using Ndns.Packet.Parser;
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net;
using System.Net.Sockets;
using System.Runtime.ExceptionServices;
using System.Threading;
using System.Threading.Tasks;

namespace Ndns.Net
{

#if NDNS_INTERNAL
    internal
#else
    public
#endif
    enum DnsClientMode
    {
        UdpWithTcpFallback = 0,
        UdpOnly = 1,
        TcpOnly = 2,
    }

    public sealed class DnsClient
    {

        private volatile int _isBusy = 0;
        private IEnumerator<ushort> _sequence;

        private bool _useRandomTransactionId;

#if NDNS_INTERNAL
        internal
#else
        public
#endif
            bool UseRandomTransactionId
        {
            get => _useRandomTransactionId;
            set
            {
                if (_useRandomTransactionId != value)
                {
                    _sequence.Dispose();
                    if (value)
                        _sequence = TransactionIdGenerator.GetRandomTransactionId().GetEnumerator();
                    else
                        _sequence = TransactionIdGenerator.GetSequentialTransactionId().GetEnumerator();
                    _useRandomTransactionId = value;
                }
            }
        }
        public IPEndPoint ServerEndPoint { get; set; }

        public bool EnableCompression { get; set; } = true;


#if NDNS_INTERNAL
        internal
#else
        public
#endif
            DnsClientMode Mode
        { get; set; } = DnsClientMode.UdpWithTcpFallback;

#if NDNS_INTERNAL
        internal
#else
    public
#endif
        int SingleUdpQueryTimeout
        { get; set; } = 2000;

#if NDNS_INTERNAL
        internal
#else
    public
#endif
        int UdpRetryCount
        { get; set; } = 3;

#if NDNS_INTERNAL
        internal
#else
    public
#endif
        int TcpQueryTimeout
        { get; set; } = 30000;

#if NDNS_INTERNAL
        internal
#else
    public
#endif
        int TotalTimeout
        { get; set; }

        private DnsPacketParser _packetParser;

        public DnsClient()
        {
            _useRandomTransactionId = false;
            _sequence = TransactionIdGenerator.GetSequentialTransactionId().GetEnumerator();
            _packetParser = new DnsPacketParser();
        }

        public DnsClient(DnsPacketParser packetParser)
        {
            _useRandomTransactionId = false;
            _sequence = TransactionIdGenerator.GetSequentialTransactionId().GetEnumerator();
            _packetParser = packetParser ?? new DnsPacketParser();
        }

#if NDNS_INTERNAL
        internal
#else
        public
#endif
        async Task<DnsPacket> SendAsync(DnsQuery query, CancellationToken cancellationToken = default(CancellationToken))
        {
            if (query == null)
                throw new ArgumentNullException(nameof(query));
            if (ServerEndPoint == null)
                throw new InvalidOperationException("ServerEndPoint must be specified.");

            if (1 == Interlocked.CompareExchange(ref _isBusy, 1, 0))
                throw new InvalidOperationException();

            try
            {
                if (TotalTimeout > 0)
                {
                    var totalTimeoutCts = new CancellationTokenSource(TotalTimeout);
                    cancellationToken = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, totalTimeoutCts.Token).Token;
                }
                return await InternalSendAsync(query, cancellationToken).ConfigureAwait(false);
            }
            finally
            {
                _isBusy = 0;
            }
        }
        private async Task<DnsPacket> InternalSendAsync(DnsQuery query, CancellationToken cancellationToken)
        {
            DnsPacket request = new DnsPacket(0, DnsPacketFlag.RecursionDesired);
            request.Queries.Add(query);
            DnsPacket response = null;
            var errors = new List<ExceptionDispatchInfo>();

            if (Mode != DnsClientMode.TcpOnly)
            {
                int udpCount = 1 + (UdpRetryCount >= 0 ? UdpRetryCount : 0);
                for (int i = 1; i <= udpCount; i++)
                {
                    try
                    {
                        var ct = cancellationToken;
                        if (SingleUdpQueryTimeout > 0)
                        {
                            var singleTimeoutCts = new CancellationTokenSource(SingleUdpQueryTimeout);
                            ct = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, singleTimeoutCts.Token).Token;
                        }
                        response = await InternalSendOverUdpAsync(request, ct).ConfigureAwait(false);
                    }
                    catch (OperationCanceledException ex)
                    {
                        errors.Add(ExceptionDispatchInfo.Capture(ex));
                        continue;
                    }
                    catch (Exception ex)
                    {
                        errors.Add(ExceptionDispatchInfo.Capture(ex));
                        break;
                    }
                    if (response != null)
                        break;
                }
            }

            if (response != null && !response.Truncated)
                return response;

            if (Mode != DnsClientMode.UdpOnly)
            {
                try
                {
                    var ct = cancellationToken;
                    if (TcpQueryTimeout > 0)
                    {
                        var singleTimeoutCts = new CancellationTokenSource(TcpQueryTimeout);
                        ct = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, singleTimeoutCts.Token).Token;
                    }
                    response = await InternalSendOverTcpAsync(request, ct).ConfigureAwait(false);
                }
                catch (Exception ex)
                {
                    errors.Add(ExceptionDispatchInfo.Capture(ex));
                }
            }

            if (response == null)
            {
                throw new AggregateException(errors.Select(edi => edi.SourceException));
            }

            return response;
        }
        private Socket _udpSocket;
        private void InitializeUdpSocket()
        {
            if (_udpSocket != null)
            {
                if (_udpSocket.AddressFamily == ServerEndPoint.AddressFamily)
                    return;
                _udpSocket.Dispose();
            }
            _udpSocket = new Socket(ServerEndPoint.AddressFamily, SocketType.Dgram, ProtocolType.Udp);
            _udpSocket.Bind(new IPEndPoint(ServerEndPoint.AddressFamily == AddressFamily.InterNetworkV6 ? IPAddress.IPv6Any : IPAddress.Any, 0));
        }
        private async Task<DnsPacket> InternalSendOverUdpAsync(DnsPacket request, CancellationToken cancellationToken)
        {
            var parser = new DnsPacketParser();

            InitializeUdpSocket();

            _sequence.MoveNext();
            request.TransactionID = _sequence.Current;

            try
            {
                //数据包
                var reqPacketBuff = request.GetBytes(true, EnableCompression);
                //接收缓冲区
                var buffer = new byte[4096];
                var taskReceive = _udpSocket.ReceiveFromAsync(new ArraySegment<byte>(buffer), SocketFlags.None, ServerEndPoint).ConfigureAwait(false);
                //发送数据包
                await _udpSocket.SendToAsync(new ArraySegment<byte>(reqPacketBuff), SocketFlags.None, ServerEndPoint).ConfigureAwait(false);
                while (true)
                {
                    //获取接收的数据
                    var received = await taskReceive;
                    //解析
                    using (var ms = new MemoryStream(buffer, 0, received.ReceivedBytes, false))
                    {
                        if (!_packetParser.TryParse(ms, out var response, out var c) && !response.Truncated)
                            throw new FormatException();
                        if (request.TransactionID == response.TransactionID)
                            return response;
                    }
                    //ID不对，继续接收
                    taskReceive = _udpSocket.ReceiveFromAsync(new ArraySegment<byte>(buffer), SocketFlags.None, ServerEndPoint).ConfigureAwait(false);
                }
            }
            catch (ObjectDisposedException)
            {
                cancellationToken.ThrowIfCancellationRequested();
                throw;
            }
        }
        private async Task<DnsPacket> InternalSendOverTcpAsync(DnsPacket request, CancellationToken cancellationToken)
        {
            _sequence.MoveNext();
            request.TransactionID = _sequence.Current;

            //准备Socket
            using (var tcpSocket = new Socket(ServerEndPoint.AddressFamily, SocketType.Stream, ProtocolType.Tcp))
            using (cancellationToken.Register(() => tcpSocket.Dispose()))
            {
                try
                {
                    DnsPacket packet;
                    //连接服务器
                    await tcpSocket.ConnectAsync(ServerEndPoint.Address, ServerEndPoint.Port).ConfigureAwait(false);
                    using (var ns = new NetworkStream(tcpSocket, false))
                    using (var ms = new MemoryStream(512))
                    {
                        //发送数据
                        int count = EnableCompression ? request.CopyCompressedTo(ms) : request.CopyTo(ms);
                        ms.Seek(0, SeekOrigin.Begin);
                        ns.WriteNetworkUInt16((ushort)count);
                        await ms.CopyToAsync(ns).ConfigureAwait(false);
                        await ns.FlushAsync().ConfigureAwait(false);
                        //接收数据
                        ms.Seek(0, SeekOrigin.Begin);
                        ms.SetLength(0);
                        while (ms.Length < 2)
                            await ns.CopyToAsync(ms).ConfigureAwait(false);
                        ms.Seek(0, SeekOrigin.Begin);
                        int expectedLength = ms.ReadNetworkUInt16() + 2;
                        ms.Seek(0, SeekOrigin.End);
                        while (ms.Length < expectedLength)
                            await ns.CopyToAsync(ms).ConfigureAwait(false);
                        ms.Seek(2, SeekOrigin.Begin);
                        if (!_packetParser.TryParse(ms, out packet, out var c) && !(packet?.Truncated ?? false))
                            throw new FormatException();
                    }
                    //关闭连接
                    tcpSocket.Shutdown(SocketShutdown.Both);
                    //返回
                    return packet;
                }
                catch (ObjectDisposedException)
                {
                    cancellationToken.ThrowIfCancellationRequested();
                    throw;
                }
            }
        }



        public static async Task<DnsPacket> SendOverUdpToAsync(DnsPacket request, IPEndPoint server, bool enableCompression = true, CancellationToken cancellationToken = default(CancellationToken))
        {
            var parser = new DnsPacketParser();
            //准备Socket
            using (var udpSocket = new Socket(server.AddressFamily, SocketType.Dgram, ProtocolType.Udp))
            using (cancellationToken.Register(() => udpSocket.Dispose()))
            {
                udpSocket.Bind(new IPEndPoint(server.AddressFamily == AddressFamily.InterNetworkV6 ? IPAddress.IPv6Any : IPAddress.Any, 0));
                try
                {
                    //接收数据
                    var buffer = new byte[4096];
                    var taskReceive = udpSocket.ReceiveFromAsync(new ArraySegment<byte>(buffer), SocketFlags.None, server).ConfigureAwait(false);
                    //发送数据包
                    var reqPacketBuff = request.GetBytes(true, enableCompression);
                    await udpSocket.SendToAsync(new ArraySegment<byte>(reqPacketBuff), SocketFlags.None, server).ConfigureAwait(false);
                    //获取接收的数据
                    var received = await taskReceive;
                    using (var ms = new MemoryStream(buffer))
                    {
                        if (!parser.TryParse(ms, out var packet, out var c) && !(packet?.Truncated ?? false))
                            throw new FormatException();
                        return packet;
                    }
                }
                catch (Exception)
                {
                    cancellationToken.ThrowIfCancellationRequested();
                    throw;
                }
            }
        }


        public static async Task<DnsPacket> SendOverTcpToAsync(DnsPacket request, IPEndPoint server, bool enableCompression = true, CancellationToken cancellationToken = default(CancellationToken))
        {
            var parser = new DnsPacketParser();
            //准备Socket
            using (var tcpSocket = new Socket(server.AddressFamily, SocketType.Stream, ProtocolType.Tcp))
            using (cancellationToken.Register(() => tcpSocket.Dispose()))
            {
                try
                {
                    DnsPacket packet;
                    //连接服务器
                    await tcpSocket.ConnectAsync(server.Address, server.Port).ConfigureAwait(false);
                    using (var ns = new NetworkStream(tcpSocket, false))
                    using (var ms = new MemoryStream(512))
                    {
                        //发送数据
                        ms.SetLength(2);
                        ms.Seek(2, SeekOrigin.Begin);
                        int count = enableCompression ? request.CopyCompressedTo(ms) : request.CopyTo(ms);
                        ms.Seek(0, SeekOrigin.Begin);
                        ms.WriteNetworkUInt16((ushort)count);
                        ms.Seek(0, SeekOrigin.Begin);
                        await ms.CopyToAsync(ns).ConfigureAwait(false);
                        //接收数据
                        ms.Seek(0, SeekOrigin.Begin);
                        ms.SetLength(0);
                        await ns.CopyBytesToAsync(ms, 2).ConfigureAwait(false);
                        ms.Seek(0, SeekOrigin.Begin);
                        int expectedLength = ms.ReadNetworkUInt16();
                        await ns.CopyBytesToAsync(ms, expectedLength).ConfigureAwait(false);
                        ms.Seek(2, SeekOrigin.Begin);
                        if (!parser.TryParse(ms, out packet, out var c) && !(packet?.Truncated ?? false))
                            throw new FormatException();
                    }
                    //关闭连接
                    tcpSocket.Shutdown(SocketShutdown.Both);
                    //返回
                    return packet;
                }
                catch (Exception)
                {
                    cancellationToken.ThrowIfCancellationRequested();
                    throw;
                }
            }
        }
    }
}
